//	Aviv Charikar	Yael Tzirulnikov	305343600	307955609

package modelTraining;

import java.io.*;
import java.util.*;
import java.util.Map.Entry;

public class HeldOutModelTraining {
	private class CountAndProbability {
		public int count;
		public double probability;
		
		public CountAndProbability(int count, double probability) {
			this.count = count;
			this.probability = probability;
		}
	}
	
	private int vocSize;
	private HashMap<String, CountAndProbability> trainingSet;
	private int countTraining;
	private double unseenProb;
	private HashMap<String, Integer> heldoutSet;
	private int countHeldout;
	private HashMap<Integer, CountSumAndEventCount> rToParamsMap;
	
	public HeldOutModelTraining(int vocSize, String devFilename, int trainSize) throws IOException {
		this.vocSize = vocSize;
		trainingSet = new HashMap<String, CountAndProbability>();
		countTraining = 0;
		heldoutSet = new HashMap<String, Integer>();
		countHeldout = 0;
		readDevFile(devFilename, trainSize);
		calcProbabilities();
	}
	
	private void readDevFile(String devFilename, int trainSize) throws IOException {
		BufferedReader reader = new BufferedReader(new FileReader(devFilename));
		String line;
		while ((line = reader.readLine()) != null) {
			// Handle only sentences
			if ((line.length() == 0) || (line.startsWith("<TRAIN")))
				continue;
			
			// Split the words in the sentence and add them to the counts
			String[] lineWords = line.split(" ");
			for (String word : lineWords) {
				if (word.length() > 0) {
					// Add the first words to the training set and the last to the held-out set
					if (countTraining < trainSize) {
						CountAndProbability countAndProbability = trainingSet.get(word);
						if (countAndProbability == null)
							trainingSet.put(word, new CountAndProbability(1, 0));
						else
							countAndProbability.count++;
						countTraining++;
					}
					else {
						Integer count = heldoutSet.get(word);
						if (count == null)
							heldoutSet.put(word, 1);
						else
							heldoutSet.put(word, count + 1);
						countHeldout++;
					}
				}
			}
		}
		reader.close();
	}
	
	private class CountSumAndEventCount {
		public int countSum;
		public int eventsCount;
		
		public CountSumAndEventCount(int countSum, int eventsCount) {
			this.countSum = countSum;
			this.eventsCount = eventsCount;
		}
	}
	
	private void calcProbabilities() {
		// Keeps the r classes of the training set and their parameters: N(r) and C_H(r)
		rToParamsMap = new HashMap<Integer, CountSumAndEventCount>();
		// Calculates the C_H(0) i.e. the size of the held-out minus all the seen-in-training events
		int unseenCountSum = countHeldout;
		// Calculates the N(0) i.e. the vocabulary size minus the seen-in-training words
		int unseenEventsCount = vocSize;
		// Iterate over all the training set words
		for (Entry<String, CountAndProbability> entry : trainingSet.entrySet()) {
			// Get the word count in the held-out
			Integer hCount = heldoutSet.get(entry.getKey());
			if (hCount == null)
				hCount = 0;
			unseenCountSum -= hCount;
			
			// Get the r class in the training set
			int r = entry.getValue().count;
			// Update the r class parameters
			CountSumAndEventCount rCounts = rToParamsMap.get(r);
			if (rCounts == null)
				rToParamsMap.put(r, new CountSumAndEventCount(hCount, 1));
			else {
				rCounts.countSum += hCount;
				rCounts.eventsCount++;
			}
			unseenEventsCount--;
		}
		rToParamsMap.put(0, new CountSumAndEventCount(unseenCountSum, unseenEventsCount));
		
		// Calculate the word probabilities
		for (Entry<String, CountAndProbability> entry : trainingSet.entrySet()) {
			// Get the r class in the training set
			int r = entry.getValue().count;
			CountSumAndEventCount rCounts = rToParamsMap.get(r);
			// Calculate the probability of the r class
			entry.getValue().probability = (double)(rCounts.countSum) / rCounts.eventsCount / countHeldout;
		}
		unseenProb = (double)unseenCountSum / unseenEventsCount / countHeldout;
	}
	
	public Set<String> getAllTrainingWords() {
		return trainingSet.keySet();
	}
	
	public int getNumOfWordsWithCount(int r) {
		CountSumAndEventCount rCounts = rToParamsMap.get(r);
		if (rCounts == null)
			return 0;
		return rCounts.eventsCount;
	}
	
	public int getTr(int r) {
		CountSumAndEventCount rCounts = rToParamsMap.get(r);
		if (rCounts == null)
			return 0;
		return rCounts.countSum;
	}
	
	public int getWordCount(String word) {
		// Return the count of the word or 0 if the word is unseen
		CountAndProbability countAndProbability = trainingSet.get(word);
		if (countAndProbability == null)
			return 0;
		return countAndProbability.count;
	}
	
	public double getWordFrequencyByCount(int r) {
		if(r == 0)
			return unseenProb * countTraining;
		else {
			CountSumAndEventCount rCounts = rToParamsMap.get(r);
			return countTraining * ((double)(rCounts.countSum) / (rCounts.eventsCount * countHeldout));
		}
	}
	
	public double getWordProbability(String word) {
		// Return the probability of the word or the unseen probability if the word is unseen
		CountAndProbability countAndProbability = trainingSet.get(word);
		if (countAndProbability == null)
			return unseenProb;
		return countAndProbability.probability;
	}
	
	public double getUnseenWordProbability() {
		return unseenProb;
	}
	
	public int getCountTraining() {
		return countTraining;
	}
	
	public int getCountHeldout() {
		return countHeldout;
	}
	
	public double calcPerplexity(BasicModelTraining testSet) {
		// Calculate the perplexity on the test set
		double perp = 0;
		for (String word : testSet.getAllWords()) {
			int wordTestCount = testSet.getWordCount(word);
			perp += wordTestCount * Math.log(getWordProbability(word)) / Math.log(2);
		}
		perp /= testSet.getCountAll();
		perp = Math.pow(2, -perp);
		return perp;
	}
}
